import os  # 导入操作系统接口，用于文件路径操作
import warnings  # 导入警告模块，用于处理警告信息

import matplotlib  # 导入matplotlib库，用于绘图
matplotlib.use('Agg')  # 设置matplotlib的后端为Agg，用于无头环境绘图
from matplotlib import pyplot as plt  # 导入pyplot模块，用于创建图表
import scipy.signal  # 导入scipy的信号处理模块，用于信号处理

import cv2  # 导入OpenCV库，用于图像处理
import shutil  # 导入shutil库，用于文件操作
import numpy as np  # 导入numpy库，用于科学计算
import tensorflow as tf  # 导入TensorFlow库

from tensorflow import keras  # 从TensorFlow中导入Keras模块
from tensorflow.keras import backend as K  # 从Keras中导入后端接口
from tqdm import tqdm  # 导入tqdm库，用于显示进度条
from .utils import cvtColor, preprocess_input, resize_image  # 从当前包的utils模块导入函数
from .utils_metrics import compute_mIoU  # 从当前包的utils_metrics模块导入计算mIoU的函数
from osgeo import gdal  # 导入GDAL库，用于地理空间数据处理

# 定义LossHistory类，用于记录训练过程中的损失值
class LossHistory(keras.callbacks.Callback):
    def __init__(self, log_dir, val_loss_flag=True):
        self.log_dir = log_dir  # 设置日志目录
        self.losses = []  # 初始化损失值列表
        self.val_loss_flag = val_loss_flag  # 是否记录验证损失
        self.losses = []  # 重复初始化，可能是代码错误，应该删除
        if self.val_loss_flag:
            self.val_loss = []  # 如果记录验证损失，初始化验证损失列表

        os.makedirs(self.log_dir)  # 创建日志目录

    # 当一个训练周期结束时调用
    def on_epoch_end(self, epoch, logs={}):
        if not os.path.exists(self.log_dir):  # 如果日志目录不存在，则创建
            os.makedirs(self.log_dir)

        self.losses.append(logs.get('loss'))  # 记录训练损失
        if self.val_loss_flag:
            self.val_loss.append(logs.get('val_loss'))  # 如果记录验证损失，记录验证损失

        # 将损失值写入文件
        with open(os.path.join(self.log_dir, "epoch_loss.txt"), 'a') as f:
            f.write(str(logs.get('loss')))
            f.write("\n")
        if self.val_loss_flag:
            with open(os.path.join(self.log_dir, "epoch_val_loss.txt"), 'a') as f:
                f.write(str(logs.get('val_loss')))
                f.write("\n")
        self.loss_plot()  # 绘制损失曲线图

    # 绘制损失曲线图的方法
    def loss_plot(self):
        iters = range(len(self.losses))

        plt.figure()
        plt.plot(iters, self.losses, 'red', linewidth=2, label='train loss')
        if self.val_loss_flag:
            plt.plot(iters, self.val_loss, 'coral', linewidth=2, label='val loss')

        try:
            if len(self.losses) < 25:
                num = 5
            else:
                num = 15

            plt.plot(iters, scipy.signal.savgol_filter(self.losses, num, 3), 'green', linestyle='--', linewidth=2, label='smooth train loss')
            if self.val_loss_flag:
                plt.plot(iters, scipy.signal.savgol_filter(self.val_loss, num, 3), '#8B4513', linestyle='--', linewidth=2, label='smooth val loss')
        except:
            pass

        plt.grid(True)
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.title('A Loss Curve')
        plt.legend(loc="upper right")

        plt.savefig(os.path.join(self.log_dir, "epoch_loss.png"))

        plt.cla()
        plt.close("all")

# 定义ExponentDecayScheduler类，用于指数衰减学习率调度
class ExponentDecayScheduler(keras.callbacks.Callback):
    def __init__(self,
                 decay_rate,
                 verbose=0):
        super(ExponentDecayScheduler, self).__init__()  # 调用父类的构造函数
        self.decay_rate = decay_rate  # 衰减率
        self.verbose = verbose  # 是否打印详细信息
        self.learning_rates = []  # 记录学习率

    def on_epoch_end(self, batch, logs=None):
        learning_rate = K.get_value(self.model.optimizer.lr) * self.decay_rate  # 计算新的学习率
        K.set_value(self.model.optimizer.lr, learning_rate)  # 更新学习率
        if self.verbose > 0:  # 如果verbose大于0，则打印学习率
            print('Setting learning rate to %s.' % (learning_rate))

# 定义EvalCallback类，用于在训练过程中评估模型性能
class EvalCallback(keras.callbacks.Callback):
    def __init__(self, model_body, input_shape, num_classes, image_ids, dataset_path, log_dir,
                 miou_out_path=".temp_miou_out", eval_flag=True, period=1):
        super(EvalCallback, self).__init__()  # 调用父类的构造函数

        self.model_body = model_body  # 模型主体
        self.input_shape = input_shape  # 输入形状
        self.num_classes = num_classes  # 类别数量
        self.image_ids = image_ids  # 图像ID列表
        self.dataset_path = dataset_path  # 数据集路径
        self.log_dir = log_dir  # 日志目录
        self.miou_out_path = miou_out_path  # mIoU输出路径
        self.eval_flag = eval_flag  # 是否执行评估
        self.period = period  # 评估周期

        self.image_ids = [image_id.split()[0] for image_id in image_ids]  # 解析图像ID
        self.mious = [0]  # 初始化mIoU列表
        self.epoches = [0]  # 初始化周期列表
        if self.eval_flag:
            with open(os.path.join(self.log_dir, "epoch_miou.txt"), 'a') as f:
                f.write(str(0))
                f.write("\n")  # 写入初始mIoU值

    @tf.function  # 使用tf.function装饰器，加速函数执行

    def get_pred(self, image_data):
        pr = self.model_body(image_data, training=False)# 获取模型预测
        return pr

    def get_miou_png(self, image):
        # ---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        # ---------------------------------------------------------#
        image = cvtColor(image)
        orininal_h = np.array(image).shape[0]
        orininal_w = np.array(image).shape[1]
        # ---------------------------------------------------------#
        #   给图像增加灰条，实现不失真的resize
        #   也可以直接resize进行识别
        # ---------------------------------------------------------#
        image_data, nw, nh = image, 256, 256
        # ---------------------------------------------------------#
        #   归一化+通道数调整到第一维度+添加上batch_size维度
        # ---------------------------------------------------------#
        image_data = np.expand_dims(preprocess_input(np.array(image_data, np.float64)), 0)

        # ---------------------------------------------------#
        #   图片传入网络进行预测
        # ---------------------------------------------------#
        pr = self.get_pred(image_data)[0].numpy()
        # --------------------------------------#
        #   将灰条部分截取掉
        # --------------------------------------#
        pr = pr[int((self.input_shape[0] - nh) // 2): int((self.input_shape[0] - nh) // 2 + nh), \
             int((self.input_shape[1] - nw) // 2): int((self.input_shape[1] - nw) // 2 + nw)]
        # --------------------------------------#
        #   进行图片的resize
        # --------------------------------------#
        pr = cv2.resize(pr, (orininal_w, orininal_h), interpolation=cv2.INTER_LINEAR)
        # ---------------------------------------------------#
        #   取出每一个像素点的种类
        # ---------------------------------------------------#
        pr = pr.argmax(axis=-1)

        image = np.uint8(pr)
        return image
    # 保存模型的代码，


def on_epoch_end(self, epoch, logs=None):
    # 将当前的训练周期加1，因为epoch是从0开始计数的
    temp_epoch = epoch + 1

    # 如果当前周期是评估周期（temp_epoch % self.period == 0），并且评估标志self.eval_flag为True
    if temp_epoch % self.period == 0 and self.eval_flag:
        # 定义真实的分割标签目录路径
        gt_dir = os.path.join(self.dataset_path, "VOC2007/SegmentationClass/")
        # 定义预测结果的目录路径
        pred_dir = os.path.join(self.miou_out_path, 'detection-results')

        # 如果预测结果目录不存在，则创建它
        if not os.path.exists(self.miou_out_path):
            os.makedirs(self.miou_out_path)
        # 如果预测结果目录不存在，则创建它
        if not os.path.exists(pred_dir):
            os.makedirs(pred_dir)

        # 打印开始获取mIoU的提示信息
        print("Get miou.")

        # 使用tqdm来显示进度条，遍历所有图像ID
        for image_id in tqdm(self.image_ids):
            # 构建图像的完整路径
            image_path = os.path.join(self.dataset_path, "VOC2007/JPEGImages/" + image_id + ".tif")
            # 使用GDAL打开图像文件
            dataset_img = gdal.Open(image_path)
            # 获取图像的宽度和高度
            width = dataset_img.RasterXSize
            height = dataset_img.RasterYSize

            # 读取图像数据，并将其从GDAL格式转换为NumPy数组
            image = dataset_img.ReadAsArray(0, 0, width, height).transpose([1, 2, 0])

            # 调用get_miou_png方法处理图像，并获取预测结果
            image = self.get_miou_png(image)
            # 将预测结果保存为PNG图像
            cv2.imwrite(os.path.join(pred_dir, image_id + ".png"), image)

        # 打印开始计算mIoU的提示信息
        print("Calculate miou.")

        # 调用compute_mIoU函数计算mIoU
        _, IoUs, _, _ = compute_mIoU(gt_dir, pred_dir, self.image_ids, self.num_classes, None)
        # 计算所有IoU的平均值，并转换为百分比
        temp_miou = np.nanmean(IoUs) * 100

        # 将当前的mIoU值和周期添加到列表中
        self.mious.append(temp_miou)
        self.epoches.append(temp_epoch)

        # 打开或创建epoch_miou.txt文件，并追加当前的mIoU值
        with open(os.path.join(self.log_dir, "epoch_miou.txt"), 'a') as f:
            f.write(str(temp_miou))
            f.write("\n")

        # 创建一个新的图表，绘制mIoU随周期变化的曲线
        plt.figure()
        plt.plot(self.epoches, self.mious, 'red', linewidth=2, label='train miou')

        # 设置图表的网格、坐标轴标签、标题和图例
        plt.grid(True)
        plt.xlabel('Epoch')
        plt.ylabel('Miou')
        plt.title('A Miou Curve')
        plt.legend(loc="upper right")

        # 保存图表为PNG图像
        plt.savefig(os.path.join(self.log_dir, "epoch_miou.png"))
        # 清除当前图表，并关闭所有图表
        plt.cla()
        plt.close("all")

        # 打印完成获取mIoU的提示信息
        print("Get miou done.")

        # 删除预测结果目录，清理文件
        shutil.rmtree(self.miou_out_path)
# 定义ModelCheckpoint类，用于保存模型检查点

class ModelCheckpoint(keras.callbacks.Callback):
    # 构造函数，初始化ModelCheckpoint类的实例
    def __init__(self, filepath, monitor='val_loss', verbose=0,
                 save_best_only=False, save_weights_only=False,
                 mode='auto', period=1):
        super(ModelCheckpoint, self).__init__()  # 调用父类构造函数
        self.monitor = monitor  # 监控的指标
        self.verbose = verbose  # 是否打印详细信息
        self.filepath = filepath  # 保存模型的文件路径
        self.save_best_only = save_best_only  # 是否仅保存最佳模型
        self.save_weights_only = save_weights_only  # 是否仅保存模型权重
        self.period = period  # 保存模型的周期
        self.epochs_since_last_save = 0  # 自上次保存以来经过的周期数

        # 设置监控指标的比较操作和最佳值
        if mode not in ['auto', 'min', 'max']:
            warnings.warn('ModelCheckpoint mode %s is unknown, '
                          'fallback to auto mode.' % (mode),
                          RuntimeWarning)
            mode = 'auto'

        if mode == 'min':
            self.monitor_op = np.less  # 如果是最小化指标，使用小于操作
            self.best = np.Inf  # 设置最佳值为无穷大
        elif mode == 'max':
            self.monitor_op = np.greater  # 如果是最大化指标，使用大于操作
            self.best = -np.Inf  # 设置最佳值为负无穷大
        else:
            if 'acc' in self.monitor or self.monitor.startswith('fmeasure'):
                self.monitor_op = np.greater  # 如果是准确率或F1分数，使用大于操作
                self.best = -np.Inf  # 设置最佳值为负无穷大
            else:
                self.monitor_op = np.less  # 否则使用小于操作
                self.best = np.Inf  # 设置最佳值为无穷大

    # 当一个训练周期结束时调用
    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}  # 如果logs为空，则设置为一个空字典
        self.epochs_since_last_save += 1  # 增加自上次保存以来的周期数

        # 如果周期数达到设定的保存周期，或者不是仅保存最佳模型
        if self.epochs_since_last_save >= self.period:
            self.epochs_since_last_save = 0  # 重置周期计数器
            filepath = self.filepath.format(epoch=epoch + 1, **logs)  # 格式化文件路径

            # 如果是仅保存最佳模型，并且当前周期的监控指标可用
            if self.save_best_only:
                current = logs.get(self.monitor)
                if current is None:
                    warnings.warn('Can save best model only with %s available, '
                                  'skipping.' % (self.monitor), RuntimeWarning)
                else:
                    # 如果当前指标比最佳值更好
                    if self.monitor_op(current, self.best):
                        if self.verbose > 0:
                            print('\nEpoch %05d: %s improved from %0.5f to %0.5f,'
                                  ' saving model to %s'
                                  % (epoch + 1, self.monitor, self.best,
                                     current, filepath))
                        self.best = current  # 更新最佳值
                        # 保存模型权重或整个模型
                        if self.save_weights_only:
                            self.model.save_weights(filepath, overwrite=True)
                        else:
                            self.model.save(filepath, overwrite=True)
                    else:
                        if self.verbose > 0:
                            print('\nEpoch %05d: %s did not improve'
                                  % (epoch + 1, self.monitor))
            else:
                # 如果不是仅保存最佳模型，直接保存模型
                if self.verbose > 0:
                    print('\nEpoch %05d: saving model to %s'
                          % (epoch + 1, filepath))
                # 保存模型权重或整个模型
                if self.save_weights_only:
                    self.model.save_weights(filepath, overwrite=True)
                else:
                    self.model.save(filepath, overwrite=True)